use std::{collections::BTreeMap, fmt::Display};

#[derive(Debug)]
pub enum Error {
    MissingApplicationName,

    ParseValueError,
    UndefinedFlag { flag: String },
    MissingStandaloneArgument { flag: String },
    LongFlagWithSingleDash,
    InvalidFlagAfterParse,

    CompressedNonBooleanShortFlags { flag: String, short: char },
}

impl Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::MissingApplicationName => {
                write!(
                    f,
                    "Missing application name, arguments array must not be empty."
                )
            }
            Error::ParseValueError => {
                write!(f, "Failed to parse flag value as type of the it.")
            }
            Error::UndefinedFlag { flag } => {
                write!(f, "Undefined flag: {flag}")
            }
            Error::MissingStandaloneArgument { flag } => {
                write!(f, "Missing standalone argument for flag: {flag}")
            }
            Error::LongFlagWithSingleDash => todo!(),
            Error::InvalidFlagAfterParse => todo!(),

            Error::CompressedNonBooleanShortFlags { flag, short } => {
                write!(f, "Type of `-{short}' is not a boolean, so `-{short}' in compressed flags `{flag}' is invalid")
            }
        }
    }
}

pub trait FlagValue: std::fmt::Debug {
    /// How to parse a from a string in command line.
    fn parse(&mut self, option: &str) -> Result<(), Error>;

    /// Is this a switch flag (with boolean argument).
    fn is_switch(&self) -> bool {
        return false;
    }

    /// Is this boolean true?.
    fn is_true(&self) -> bool {
        return false;
    }
}

impl<T> FlagValue for Option<T>
where
    T: FlagValue + Default,
{
    fn parse(&mut self, option: &str) -> Result<(), Error> {
        match self {
            Some(x) => x.parse(option)?,
            None => {
                let mut x = T::default();
                x.parse(option)?;
                *self = Some(x)
            }
        }

        Ok(())
    }

    fn is_switch(&self) -> bool {
        return self.as_ref().map(|x| x.is_switch()).unwrap_or(false);
    }

    fn is_true(&self) -> bool {
        return self.as_ref().map(|x| x.is_switch()).unwrap_or(false);
    }
}

impl<T> FlagValue for Vec<T>
where
    T: FlagValue + Default,
{
    fn parse(&mut self, option: &str) -> Result<(), Error> {
        self.clear();
        for each in option.split(',') {
            let mut x = T::default();
            x.parse(each)?;
            self.push(x);
        }
        Ok(())
    }
}

impl FlagValue for bool {
    fn parse(&mut self, option: &str) -> Result<(), Error> {
        match option {
            "t" | "true" | "1" | "on" | "yes" | "y" => *self = true,
            "f" | "false" | "0" | "off" | "no" | "n" => *self = false,
            _ => return Err(Error::ParseValueError),
        }
        Ok(())
    }

    fn is_switch(&self) -> bool {
        return true;
    }

    fn is_true(&self) -> bool {
        return *self;
    }
}

macro_rules! impl_simple_flag_value {
    ($type: ty) => {
        impl FlagValue for $type {
            fn parse(&mut self, option: &str) -> Result<(), Error> {
                *self = option.parse().map_err(|_| Error::ParseValueError)?;
                Ok(())
            }
        }
    };
}

impl_simple_flag_value!(i8);
impl_simple_flag_value!(u8);
impl_simple_flag_value!(i16);
impl_simple_flag_value!(u16);
impl_simple_flag_value!(i32);
impl_simple_flag_value!(u32);
impl_simple_flag_value!(i64);
impl_simple_flag_value!(u64);
impl_simple_flag_value!(f32);
impl_simple_flag_value!(f64);

impl FlagValue for String {
    fn parse(&mut self, option: &str) -> Result<(), Error> {
        *self = option.to_string();
        Ok(())
    }
}

#[derive(Debug, Clone)]
pub struct FlagConfig {
    pub name: String,
    pub short: Option<char>,
    pub usage: String,
    pub default_value: String,
}

/// How to handle parse error.
#[derive(Clone, Copy)]
pub enum ErrorHandler {
    Exit,
    Return,
}

/// A set of config of command line flags.
pub struct FlagSet<'t> {
    flag_configs: Vec<FlagConfig>,
    flag_values: Vec<Box<&'t mut dyn FlagValue>>,

    name_map: BTreeMap<String, usize>,
    short_name_map: BTreeMap<char, usize>,

    application: String,
    error_handler: ErrorHandler,
}

/// Rest things useful after parse in FlagSet.
#[derive(Debug, Clone)]
pub struct ParsedArgs {
    /// First string of the command line arguments.
    pub application: String,

    /// Configs of flags, excluding values.
    pub flag_configs: Vec<FlagConfig>,

    /// Command line arguments not starts with '-' or after '--'.
    pub positionals: Vec<String>,
}

impl ParsedArgs {
    pub fn write_usage<W: std::io::Write>(&self, w: &mut W) -> std::io::Result<()> {
        writeln!(w, "Usage: {} [flags]... [arguments]...", self.application)?;
        writeln!(w)?;

        writeln!(w, "Flags:")?;

        for flag_config in (&self.flag_configs).iter() {
            if let Some(short) = flag_config.short {
                write!(w, "\t-{},--{}", short, flag_config.name)?;
            } else {
                write!(w, "\t--{}", flag_config.name)?;
            }
            writeln!(w, " (default: {})", flag_config.default_value)?;

            for line in flag_config.usage.split('\n') {
                writeln!(w, "\t\t{}", line)?;
            }
        }

        Ok(())
    }
}

impl<'t> FlagSet<'t> {
    pub fn new() -> FlagSet<'t> {
        FlagSet {
            flag_values: Vec::new(),
            flag_configs: Vec::new(),

            name_map: BTreeMap::new(),
            short_name_map: BTreeMap::new(),
            application: "<application name>".to_string(),
            error_handler: ErrorHandler::Exit,
        }
    }

    pub fn error_handler(&mut self, err: ErrorHandler) -> &mut Self {
        self.error_handler = err;
        return self;
    }

    pub fn flag<Name: ToString, Usage: ToString, V: FlagValue>(
        &mut self,
        name: Name,
        short: Option<char>,
        v: &'t mut V,
        usage: Usage,
    ) -> &mut Self {
        self.name_map
            .insert(name.to_string(), self.flag_values.len());

        if let Some(short) = short {
            self.short_name_map.insert(short, self.flag_values.len());
        }

        self.flag_configs.push(FlagConfig {
            name: name.to_string(),
            short,
            usage: usage.to_string(),
            default_value: format!("{:?}", v),
        });

        self.flag_values.push(Box::new(v));
        return self;
    }

    pub fn parse_args(self) -> Result<ParsedArgs, Error> {
        self.parse(&std::env::args().collect::<Vec<_>>())
    }

    pub fn parse(mut self, args: &[String]) -> Result<ParsedArgs, Error> {
        let positionals = match self.parse_inner(args) {
            Ok(ok) => ok,
            Err(Error::MissingApplicationName) => return Err(Error::MissingApplicationName),
            Err(err) => {
                if matches!(self.error_handler, ErrorHandler::Exit) {
                    let args = ParsedArgs {
                        flag_configs: self.flag_configs,
                        application: self.application,
                        positionals: vec![],
                    };

                    eprintln!("Error: {}\n", err);
                    let _ = args.write_usage(&mut std::io::stderr());
                    std::process::exit(2);
                }
                return Err(err);
            }
        };

        let parsed = ParsedArgs {
            flag_configs: self.flag_configs,
            application: self.application,
            positionals,
        };

        Ok(parsed)
    }

    fn parse_inner(&mut self, args: &[String]) -> Result<Vec<String>, Error> {
        let Some((app, mut args)) = args.split_first() else {
            return Err(Error::MissingApplicationName);
        };

        self.application = app.to_owned();

        let mut pos = vec![];

        while let Some((arg, rest)) = args.split_first() {
            let Some(arg) = arg.strip_prefix('-') else {
                pos.push(arg.to_owned());
                args = rest;
                continue;
            };

            if arg == "-" {
                pos.extend(rest.to_vec());
                break;
            }

            match arg.split_once('=') {
                Some((name, opt)) => {
                    self.parse_with_eq(name, opt)?;
                    args = rest;
                }
                None => {
                    args = self.parse_without_eq(arg, rest)?;
                }
            }
        }

        Ok(pos)
    }

    // "-abc" or "--name" or "--name opt"
    fn parse_without_eq<'s>(
        &mut self,
        option: &str,
        mut args: &'s [String],
    ) -> Result<&'s [String], Error> {
        if let Some(long) = option.strip_prefix('-') {
            let Some(flag) = self.name_map.get_mut(long) else {
                return Err(Error::UndefinedFlag {
                    flag: format!("--{}", long),
                });
            };

            let flag = &mut self.flag_values[*flag];

            if flag.is_switch() {
                flag.parse("true")?;
            } else {
                let Some((next, rest)) = args.split_first() else {
                    return Err(Error::MissingStandaloneArgument {
                        flag: format!("--{long}"),
                    });
                };

                flag.parse(next)?;
                args = rest;
            }
        } else {
            let only_one = option.len() == 1;

            for short in option.chars() {
                let Some(flag) = self.short_name_map.get_mut(&short) else {
                    return Err(Error::UndefinedFlag {
                        flag: format!("-{}", short),
                    });
                };

                let flag = &mut self.flag_values[*flag];

                if !flag.is_switch() {
                    if only_one {
                        return Err(Error::MissingStandaloneArgument {
                            flag: format!("-{short}"),
                        });
                    }

                    return Err(Error::CompressedNonBooleanShortFlags {
                        flag: format!("-{}", option),
                        short,
                    });
                }

                flag.parse("true")?;
            }
        }

        Ok(args)
    }

    // --long=opt or -c=opt
    fn parse_with_eq(&mut self, name: &str, opt: &str) -> Result<(), Error> {
        let flag = if let Some(long) = name.strip_prefix('-') {
            // --long=opt

            let Some(flag) = self.name_map.get_mut(long) else {
                return Err(Error::UndefinedFlag {
                    flag: format!("--{}", long),
                });
            };

            *flag
        } else {
            // -d=opt

            if name.len() != 1 {
                return Err(Error::LongFlagWithSingleDash);
            }

            let short = name.chars().next().unwrap();

            let Some(flag) = self.short_name_map.get_mut(&short) else {
                return Err(Error::UndefinedFlag {
                    flag: format!("-{}", short),
                });
            };

            *flag
        };

        self.flag_values[flag].parse(opt)?;

        Ok(())
    }
}

#[test]
fn test_flag() {
    #[derive(Debug, Default)]
    struct TestCase {
        args: Vec<String>,
        create: bool,
        gzip: bool,
        strip_components: i64,
        pos_args: Vec<String>,
    }

    let test_cases = [
        TestCase {
            args: vec!["tar".to_string(), "--create=true".to_string()],
            create: true,
            ..Default::default()
        },
        TestCase {
            args: vec!["tar".to_string(), "--create=false".to_string()],
            ..Default::default()
        },
        TestCase {
            args: vec!["tar".to_string(), "-cz".to_string()],
            create: true,
            gzip: true,
            ..Default::default()
        },
        TestCase {
            args: vec!["tar".to_string(), "--strip-components=1".to_string()],
            strip_components: 1,
            ..Default::default()
        },
        TestCase {
            args: vec!["tar".to_string(), "--strip-components=-10".to_string()],
            strip_components: -10,
            ..Default::default()
        },
        TestCase {
            args: vec![
                "tar".to_string(),
                "--strip-components=-10".to_string(),
                "--strip-components=10".to_string(),
            ],
            strip_components: 10,
            ..Default::default()
        },
        TestCase {
            args: vec![
                "tar".to_string(),
                "--strip-components".to_string(),
                "10".to_string(),
            ],
            strip_components: 10,
            ..Default::default()
        },
        TestCase {
            args: vec![
                "tar".to_string(),
                "input1.txt".to_string(),
                "input2.txt".to_string(),
                "-c".to_string(),
                "input3.txt".to_string(),
                "--".to_string(),
                "-z".to_string(),
            ],
            create: true,
            pos_args: vec![
                "input1.txt".to_string(),
                "input2.txt".to_string(),
                "input3.txt".to_string(),
                "-z".to_string(),
            ],
            ..Default::default()
        },
    ];

    for test_case in test_cases {
        let mut create = false;
        let mut gzip = false;
        let mut strip_comp = 0;

        let mut fs = FlagSet::new();

        fs.flag(
            "create",
            Some('c'),
            &mut create,
            concat!(
                "Create a new archive. Arguments supply the names of the files\n",
                "to be archived. Directories are archived recursively, unless\n",
                "the `--no-recursion' option is given."
            ),
        );

        fs.flag(
            "gzip",
            Some('z'),
            &mut gzip,
            "Filter the archive through gzip(1).",
        );

        fs.flag(
            "strip-components",
            None,
            &mut strip_comp,
            "Strip some leading components from file names on extraction.",
        );

        let result = fs.parse(&test_case.args).unwrap();

        assert_eq!(create, test_case.create);
        assert_eq!(gzip, test_case.gzip);
        assert_eq!(strip_comp, test_case.strip_components);
        assert_eq!(result.positionals, test_case.pos_args);
    }
}
